import torch
import random
class AudioAugmenter:
    def __init__(self, sample_rate=32000, enable_time_mask=True, 
                 enable_freq_mask=True, enable_gaussian_noise=True):
        self.sample_rate = sample_rate
        self.enable_time_mask = enable_time_mask
        self.enable_freq_mask = enable_freq_mask
        self.enable_gaussian_noise = enable_gaussian_noise
        
    def __call__(self, waveform):
        # 时域增强
        if self.enable_time_mask:
            waveform = self.time_mask(waveform)
        
        # 频域增强（在频谱图空间）
        # 注意：实际在forward中应用
        
        # 加性噪声
        if self.enable_gaussian_noise:
            waveform = waveform + torch.randn_like(waveform) * 0.01
        
        return waveform
    
    def time_mask(self, x):
        """时间掩码增强"""
        x_len = x.shape[-1]
        mask_len = random.randint(int(0.05*x_len), int(0.2*x_len))
        mask_pos = random.randint(0, x_len - mask_len)
        x[..., mask_pos:mask_pos+mask_len] = 0
        return x
    
    def freq_mask(self, spec):
        """频率掩码增强（在forward中调用）"""
        if not self.enable_freq_mask:
            return spec
        
        freq_len = spec.shape[2]
        mask_len = random.randint(int(0.05*freq_len), int(0.2*freq_len))
        mask_pos = random.randint(0, freq_len - mask_len)
        spec[..., mask_pos:mask_pos+mask_len, :] = 0
        return spec